/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#include "mstep.h"
#include "qfunc.h"

#include <assert.h>
#include <gsl/gsl_multimin.h>
#include <stdlib.h>
#include <string.h>

static void
to_x  (const double*    kappa,
       const double*    logalpha,
       unsigned int     n_workers,
       unsigned int     n_labels,
       const double*    tau,
       unsigned int     n_images,
       const double*    gamma,
       const double*    rho,
       const double*    loglambda,
       gsl_vector*      x)
{ 
  size_t off = 0;

  x->data[off] = *kappa;
  ++off;
  memcpy (x->data + off, logalpha, n_images * sizeof (*logalpha));
  off += n_images;

  if (n_labels > 0)
    {
      memcpy (x->data + off, tau, n_workers * (n_labels - 1) * sizeof (*tau));
      off += n_workers * (n_labels - 1);
      memcpy (x->data + off, gamma, (n_labels - 1) * sizeof (*gamma));
      off += (n_labels - 1);
    }

  x->data[off] = *rho;
  ++off;
  memcpy (x->data + off, loglambda, n_workers * (sizeof (*loglambda)));
}

static void
from_x (double*           kappa,
        double*           logalpha,
        unsigned int      n_workers,
        unsigned int      n_labels,
        double*           tau,
        unsigned int      n_images,
        double*           gamma,
        double*           rho,
        double*           loglambda,
        const gsl_vector* x)
{ 
  size_t off = 0;

  *kappa = x->data[off];
  ++off;
  memcpy (logalpha, x->data + off, n_images * sizeof (logalpha[0]));
  off += n_images;

  if (n_labels > 0)
    {
      memcpy (tau, x->data + off, n_workers * (n_labels - 1) * sizeof (*tau));
      off += n_workers * (n_labels - 1);
      memcpy (gamma, x->data + off, (n_labels - 1) * sizeof (*gamma));
      off += (n_labels - 1);
    }

  *rho = x->data[off];
  ++off;
  memcpy (loglambda, x->data + off, n_workers * (sizeof (*loglambda)));
}

typedef struct _QfuncParams QfuncParams;
struct _QfuncParams
{ 
  unsigned int                          n_workers;
  unsigned int                          n_labels;
  unsigned int                          n_images;
  const double*                         pz;
  const double*                         priorz;
  const Rating*                         ratings;
  unsigned int                          n_ratings;
  const Distribution*                   priorkappa;
  const ParameterizedDistribution*      priorlogalpha;
  const ParameterizedDistribution*      priortau;
  Distribution* const*                  priorgamma;
  const Distribution*                   priorrho;
  const ParameterizedDistribution*      priorloglambda;
  double                                hq;
};

static double
qfunc_f (const gsl_vector* x,
         void*             params)
{
  QfuncParams *qfunc_params = (QfuncParams *) params;
  unsigned int n_workers = qfunc_params->n_workers;
  unsigned int n_images = qfunc_params->n_images;
  unsigned int n_labels = qfunc_params->n_labels;

  return -qfunc (x->data[0],
                 x->data + 1,
                 n_workers,
                 n_labels,
                 x->data + 1 + n_images,
                 n_images,
                 x->data + 1 + n_images + n_workers * (n_labels - 1),
                 x->data[1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1)],
                 x->data + 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1) + 1,
                 qfunc_params->pz,
                 qfunc_params->ratings,
                 qfunc_params->n_ratings,
                 qfunc_params->priorkappa,
                 qfunc_params->priorlogalpha,
                 qfunc_params->priortau,
                 qfunc_params->priorgamma,
                 qfunc_params->priorrho,
                 qfunc_params->priorloglambda,
                 qfunc_params->hq);
}

static void
qfunc_df (const gsl_vector* x,
          void*             params,
          gsl_vector*       g)
{
  QfuncParams *qfunc_params = (QfuncParams *) params;
  unsigned int n_workers = qfunc_params->n_workers;
  unsigned int n_labels = qfunc_params->n_labels;
  unsigned int n_images = qfunc_params->n_images;
  unsigned int n_dims = 
    1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1) + 1 + n_workers;
    
  dqfunc (g->data,
          g->data + 1,
          g->data + 1 + n_images,
          g->data + 1 + n_images + n_workers * (n_labels - 1),
          g->data + 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1),
          g->data + 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1) + 1,
          x->data[0],
          x->data + 1,
          n_workers,
          n_labels,
          x->data + 1 + n_images,
          n_images,
          x->data + 1 + n_images + n_workers * (n_labels - 1),
          x->data[1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1)],
          x->data + 1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1) + 1,
          qfunc_params->pz,
          qfunc_params->ratings,
          qfunc_params->n_ratings,
          qfunc_params->priorkappa,
          qfunc_params->priorlogalpha,
          qfunc_params->priortau,
          qfunc_params->priorgamma,
          qfunc_params->priorrho,
          qfunc_params->priorloglambda);

  for (unsigned int i = 0; i < n_dims; ++i)
    {
      g->data[i] *= -1;
    }
}

static void
qfunc_fdf (const gsl_vector* x,
           void*             params,
           double*           f,
           gsl_vector*       g)
{ 
  *f = qfunc_f (x, params);
  qfunc_df (x, params, g);
}

double
mstep (double*                          kappa,
       double*                          logalpha,
       unsigned int                     n_workers,
       unsigned int                     n_labels,
       double*                          tau,
       unsigned int                     n_images,
       double*                          gamma,
       double*                          rho,
       double*                          loglambda,
       const double*                    pz,
       const double*                    priorz,
       const Rating*                    ratings,
       unsigned int                     n_ratings,
       const Distribution*              priorkappa,
       const ParameterizedDistribution* priorlogalpha,
       const ParameterizedDistribution* priortau,
       Distribution* const*             priorgamma,
       const Distribution*              priorrho,
       const ParameterizedDistribution* priorloglambda,
       double                           hq,
       unsigned int                     maxiter,
       double                           tol)
{
  size_t n_dims = 
    1 + n_images + n_workers * (n_labels - 1) + (n_labels - 1) + 1 + n_workers;
  gsl_multimin_fdfminimizer* s =
    gsl_multimin_fdfminimizer_alloc (gsl_multimin_fdfminimizer_conjugate_pr,
                                     n_dims);
  assert (s != NULL);
  gsl_vector* x = gsl_vector_alloc (n_dims);
  assert (x != NULL);

  QfuncParams qfunc_params = { .n_workers = n_workers,
                               .n_labels = n_labels,
                               .n_images = n_images,
                               .pz = pz,
                               .priorz = priorz,
                               .ratings = ratings,
                               .n_ratings = n_ratings,
                               .priorkappa = priorkappa,
                               .priorlogalpha = priorlogalpha,
                               .priortau = priortau,
                               .priorgamma = priorgamma,
                               .priorrho = priorrho,
                               .priorloglambda = priorloglambda,
                               .hq = hq };

  gsl_multimin_function_fdf my_func = { .f = qfunc_f,
                                        .df = qfunc_df,
                                        .fdf = qfunc_fdf,
                                        .n = n_dims,
                                        .params = &qfunc_params };

  to_x (kappa, logalpha, n_workers, n_labels, tau, n_images, gamma, rho, loglambda, x);
  gsl_multimin_fdfminimizer_set (s, &my_func, x, 0.01, 10 * tol);

  double lastf = 0;
  for (unsigned int iter = 0;
       ! gsl_multimin_fdfminimizer_iterate (s) &&
       gsl_multimin_test_gradient (s->gradient, tol) == GSL_CONTINUE &&
       (iter == 0 || lastf - s->f > 10 * tol) &&
       iter < maxiter;
       ++iter,
       lastf = s->f)
    {
    }

  from_x (kappa, logalpha, n_workers, n_labels, tau, n_images, gamma, rho, loglambda, s->x);

  double qmax = -s->f;

  gsl_vector_free (x);
  gsl_multimin_fdfminimizer_free (s);

  return qmax;
}
